import sys
import json
import numpy as np
import os
from pdf2image import convert_from_path
from paddleocr import PPStructureV3
from PIL import Image, ImageDraw, ImageFont
from bs4 import BeautifulSoup
import uuid
from pathlib import Path
import export_to_word, export_to_excel



def eprint(*args, **kwargs):
    """打印到 stderr"""
    print(*args, file=sys.stderr, **kwargs)


def pdf_to_images(pdf_path, dpi=300):
    """将 PDF 转为图片列表"""
    eprint(f"Converting PDF to images: {pdf_path}")
    
    # poppler 路径配置（如果已添加到系统PATH则可以忽略此配置）
    poppler_path = None
    
    # Windows用户：如果poppler未在PATH中，请指定完整路径
    # 例如：poppler_path = r'C:\poppler\Library\bin'
    if os.name == 'nt':  # Windows系统
        # 尝试常见的poppler安装路径
        possible_paths = [
            r'C:\poppler\Library\bin',
            r'C:\Program Files\poppler\Library\bin',
            r'D:\poppler\Library\bin'
        ]
        for path in possible_paths:
            if os.path.exists(path):
                poppler_path = path
                eprint(f"Found poppler at: {poppler_path}")
                break
    
    # 转换PDF为图片
    if poppler_path:
        images = convert_from_path(pdf_path, dpi=dpi, poppler_path=poppler_path)
    else:
        images = convert_from_path(pdf_path, dpi=dpi)
    
    eprint(f"Total pages: {len(images)}")
    return images


def clean_html(html_text):
    """清理 HTML，保留纯文本"""
    soup = BeautifulSoup(html_text, "html.parser")
    return soup.get_text(separator=" ", strip=True)


# 全局模型实例（只初始化一次，避免重复下载）
_ppstructure_instance = None

def get_ppstructure():
    """获取PP-Structure实例（单例模式）"""
    global _ppstructure_instance
    if _ppstructure_instance is None:
        eprint("Initializing PP-Structure (first time)...")
        _ppstructure_instance = PPStructureV3(lang='ch')
        eprint("PP-Structure initialized successfully!")
    return _ppstructure_instance


def structure_and_visualize(images, base_output_dir="output_structure"):
    """使用 PP-Structure 进行结构化识别，并在原图上可视化结果"""
    ppstructure = get_ppstructure()  # 使用单例实例，避免重复初始化
    os.makedirs(base_output_dir, exist_ok=True)

    json_paths = []
    vis_paths = []

    unique_id = str(uuid.uuid4())
    output_dir = os.path.join(base_output_dir, unique_id)
    os.makedirs(output_dir, exist_ok=True)

    font_path = "C:/Windows/Fonts/simhei.ttf"
    font = ImageFont.truetype(font_path, 24)

    for idx, image in enumerate(images, 1):
        image_np = np.array(image)
        structure_res = ppstructure.predict(image_np)

        for res in structure_res:
            page_save_path = os.path.join(output_dir, f"page{idx}.json")
            res.save_to_json(save_path=page_save_path)
            json_paths.append(page_save_path)

        # 可视化
        img = image.copy()
        draw = ImageDraw.Draw(img)

        with open(page_save_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        for block in data.get("parsing_res_list", []):
            label = block.get("block_label", "")
            content = clean_html(block.get("block_content", ""))
            bbox = block.get("block_bbox", [])
            if len(bbox) == 4:
                x1, y1, x2, y2 = map(int, bbox)
                draw.rectangle([x1, y1, x2, y2], outline=(0, 255, 0), width=3)
                draw.text((x1, max(0, y1-30)), label, fill=(255, 0, 0), font=font)
                draw.text((x1, y2+5), content[:100]+"..." if len(content) > 100 else content, fill=(0, 0, 255), font=font)

        vis_save_path = os.path.join(output_dir, f"page{idx}_visual.png")
        img.save(vis_save_path)
        vis_paths.append(vis_save_path)

    return json_paths, vis_paths, output_dir


def main():
    if len(sys.argv) < 3:
        eprint("Usage: python pdf_parser.py <PDF_PATH> <output_type: word|excel>")
        sys.exit(1)

    pdf_path = sys.argv[1]
    output_type = sys.argv[2].lower()

    images = pdf_to_images(pdf_path)
    json_paths, vis_paths, output_dir = structure_and_visualize(images)

    # 返回的基本结果 JSON
    result = {"json_paths": json_paths, "visual_paths": vis_paths, "output_dir": output_dir}

    # 根据前端传入参数生成对应文件
    if output_type == "word":
        word_path = os.path.join(output_dir, "result.docx")
        export_to_word.export_to_word(json_paths, word_path)
        result["word_path"] = word_path
    elif output_type == "excel":
        excel_path = os.path.join(output_dir, "result.xlsx")
        export_to_excel.export_pdf_content_to_excel(json_paths, excel_path)
        result["excel_path"] = excel_path
    else:
        eprint(f"未知 output_type: {output_type}")

    # 输出 JSON 给前端
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()
